import numpy as np
import quimb as qu
import quimb.tensor as qtn
from quimb.tensor import Tensor, TensorNetwork

Z_mea_o1 = np.array([0, 0, 0, np.sqrt(2)])
rhos_map = {
	'z+': np.array([1, 0]), 'z-': np.array([0, 1]),
}
meas_map = {
	'z+': np.array([1, 0, 0, 1]) / np.sqrt(2), 'z-': np.array([1, 0, 0, -1]) / np.sqrt(2),
}
"""Setting"""
cached = True
N_qubit, N_set, N_shot, step = 10, 1, 4000, 7
J, h, dt = 0.5236, 1.0, 0.05
max_bond, cutoff = 400, 1e-10

"""simulation"""
dual_o2 = np.array([
	[1, 1],
	[0, 0],
	[0, 0],
	[1, -1]
]) / np.sqrt(2)
site_tags = list(map("I{}".format, range(N_qubit)))

import pathlib

"""Error Mitigation"""
import tqdm
import json

# Zp_o1 = np.array([1, 0, 0, 1]) / np.sqrt(2)
dual_mpo = TensorNetwork([Tensor(dual_o2, inds=(f"b{_}", f"c{_}"), tags=[f"I{_}"]) for _ in range(N_qubit)],
						 virtual=True)
Z_meas_mps = TensorNetwork([Tensor(Z_mea_o1, inds=(f"k{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)

"""Data Load"""
# data_filename = f"data/Outcomes_reduced, N_qubit={N_qubit}, N_set={N_set}, N_shot={N_shot}, step={step}, max_bond={max_bond}.json"
data_filename = f"data/Outcomes_reduced, N_qubit={N_qubit}, N_set={N_set}, N_shot={N_shot}, step={step}, max_bond=800.json"
file = pathlib.Path(data_filename)
if not file.exists(): raise FileNotFoundError(f"Data File '{data_filename}' not Exist!")
with open(file, 'r') as f:
	print(f"Loading Measurement data from File '{data_filename}' ...")
	data_dict = json.load(f)
	print(f"Loading done.")
meas_datas = data_dict['datas']
"""TEM Load"""
TEM_trotter_step_filename = f"cache/TEM_trotter_step, step={step}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
if not pathlib.Path(TEM_trotter_step_filename).exists(): raise FileNotFoundError
TEM_trotter_step = qu.load_from_disk(TEM_trotter_step_filename)
TEM_mpo = dual_mpo | TEM_trotter_step

I_o1 = np.array([np.sqrt(2), 0, 0, 0])
quasi_outcomes = []
N_test = 1000

for i in range(N_test):
	pbar = tqdm.trange(N_set, desc=f'N_test {i}, meas {i + 1}/{N_test}', leave=True)
	for j in range(N_set):
		jth_data = meas_datas[j]
		meas_setting = jth_data['meas']
		outcomes = jth_data['outcomes']

		meas = TensorNetwork([Tensor(I_o1, inds=(f"k{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)
		quasi_outcome, quasi_prob = [], 1.0
		for k in range(N_qubit):
			meas.pop_tensor(k)
			kth_mea = meas_setting[k]
			kth_meas_p = Tensor(meas_map[kth_mea + '+'], inds=(f"k{k}",), tags=[f"I{k}"])
			kth_meas_m = Tensor(meas_map[kth_mea + '-'], inds=(f"k{k}",), tags=[f"I{k}"])
			meas_p = meas | kth_meas_p
			meas_m = meas | kth_meas_m

			prob_ps, prob_ms = [], []
			for l in range(N_qubit):
				outcome = outcomes[l]
				quasi_rho_mps = TensorNetwork(
						[Tensor(rhos_map[outcome[_]], inds=(f"c{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)],
						virtual=True
				)
				TEM_rho_mps = quasi_rho_mps | TEM_mpo
				prob_p = np.real((TEM_rho_mps | meas_p).contract())
				prob_m = np.real((TEM_rho_mps | meas_m).contract())
				prob_ps.append(prob_p)
				prob_ms.append(prob_m)
			quasi_prob_p = np.mean(prob_ps)
			quasi_prob_m = np.mean(prob_ms)
			print(
					f"quasi_prob_p={quasi_prob_p}, quasi_prob_m={quasi_prob_m}, {(quasi_prob_p + quasi_prob_m) / quasi_prob}")
			if np.random.uniform(0, quasi_prob_p + quasi_prob_m) < quasi_prob_p:
				meas |= kth_meas_p
				quasi_prob = quasi_prob_p
				quasi_outcome.append(kth_mea + '+')
			else:
				meas |= kth_meas_m
				quasi_prob = quasi_prob_m
				quasi_outcome.append(kth_mea + '-')
		quasi_outcomes.append(quasi_outcome)
print(quasi_outcomes)

# for i in range(N_set):
# 	ith_data = meas_datas[i]
#
# 	mea_setting = ith_data['meas']
# 	outcomes = ith_data['outcomes']
# 	pbar = tqdm.trange(N_test, desc=f'N_set {i}, TEM {i + 1}/{N_set}', leave=True)
# 	# for l in range(N_test):
# 	for l in pbar:
# 		meas = TensorNetwork([Tensor(I_o1, inds=(f"k{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)
# 		quasi_outcome, quasi_prob = [], 1.0
# 		for j in range(N_qubit):
# 			meas.pop_tensor(j)
# 			jth_meas_p = Tensor(meas_map[mea_setting + '+'], inds=(f"k{j}",), tags=[f"I{j}"])
# 			jth_meas_m = Tensor(meas_map[mea_setting + '-'], inds=(f"k{j}",), tags=[f"I{j}"])
# 			meas_p = meas | jth_meas_p
# 			meas_m = meas | jth_meas_m
#
# 			prob_ps, prob_ms = [], []
# 			for k in range(N_shot):
# 				outcome = outcomes[k]
# 				quasi_rho_mps = TensorNetwork(
# 					[Tensor(meas_map[outcome[_]], inds=(f"c{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)],
# 					virtual=True
# 			)
# 				TEM_rho_mps = quasi_rho_mps | TEM_mpo
# 				prob_p = np.real((TEM_rho_mps | meas_p).contract())
# 				prob_m = np.real((TEM_rho_mps | meas_m).contract())
# 				prob_ps.append(prob_p)
# 				prob_ms.append(prob_m)
# 			quasi_prob_p = np.mean(prob_ps)
# 			quasi_prob_m = np.mean(prob_ms)
# 			print(f"quasi_prob_p={quasi_prob_p}, quasi_prob_m={quasi_prob_m}, {(quasi_prob_p + quasi_prob_m) / quasi_prob}")
# 			if np.random.uniform(0, quasi_prob_p + quasi_prob_m) < quasi_prob_p:
# 				meas |= jth_meas_p
# 				quasi_prob = quasi_prob_p
# 				quasi_outcome.append(mea_setting + '+')
# 			else:
# 				meas |= jth_meas_m
# 				quasi_prob = quasi_prob_m
# 				quasi_outcome.append(mea_setting + '-')
# 		quasi_outcomes.append(quasi_outcome)
#
# print(quasi_outcomes)
# import matplotlib.pyplot as plt
